"""使用jwt实现加密传输数据,多用于网站保持登录功能上"""
import base64
import json
import time
import copy
import hmac


class Jwt:
    """
    ### 使用JWT生成加密字节串或者解密字节串, 主要用于,APP保持登录功能, 解决浏览器http无状态码的问题
    """

    # %% base64加密

    @staticmethod
    def encode(user_name, key, exp=300):
        """
        基于jwt的规范,生成token的加密字符串

        :param user_name: [用户名, 字典形式传入]
        :type user_name: [dist(字典)]
        :param key: [加密钥匙(不可泄露)]
        :type key: [字符串或者字节串]
        :param exp: [生效周期,单位秒], defaults to 300s
        :type exp: int, 可选的
        :raises TypeError: [exp只能是int或者str类型]
        :return: [返回jwt规范的加密字符串]
        :rtype: [str]
        """
        # 1.jwt的header, 这是官方固定的形式
        header = {'typ': 'JWT', 'alg': 'HS256'}

        # 转成js能识别的JSON
        # separator格式化json,参数1是键值对之间的分隔符, 参数2是键值对的分隔符, sort_keys开启有序性, 这个目的是保证有序性,以免造成密文雪崩
        header_json = json.dumps(header, separators=(',', ':'), sort_keys=True)

        # 生成bs64加密, base64原生加密会有=,在加密本意就是占位符,在程序中没用, 网络传输中也浪费资源
        # 所以需要重写base64加密去掉=, 传输回来验证时,需要在decode()方法之前把=补回来
        # header_bs = base64.urlsafe_b64encode(header_json.encode()) # 此方法没有去=号
        header_bs = Jwt.ba64encode(header_json.encode())  # 重写方法,去掉=号

        # 2.jwt的payload
        # 谨防传过来的是字典或者复杂结构, 进行拷贝一份出来单独解析
        paload_ = copy.deepcopy(user_name)
        # 判断传入的类型
        if not isinstance(exp, int) and not isinstance(exp, str):
            raise TypeError(f"{exp} 的值只能是int或者str类型!")

        paload_['exp'] = time.time() + int(exp)
        paload_josn = json.dumps(
            paload_, separators=(',', ':'), sort_keys=True)
        # 进行bs64加密
        paload_bs = Jwt.ba64encode(paload_josn.encode())

        # 3.jwt的sign, 将前两次加一起再一次加密
        if isinstance(key, str):  # 判断key是不是字符串,是则转2进制
            key = key.encode()
        hm = hmac.new(key, header_bs + b'.' +
                      paload_bs + b'.', digestmod='SHA256')
        # 加密
        sign_bs = Jwt.ba64encode(hm.digest())  # digest函数返回此哈希对象的哈希值。

        # 返回
        return header_bs + b'.' + paload_bs + b'.' + sign_bs

    @staticmethod
    def ba64encode(content):
        """
        ### 生成bs64加密, base64原生加密会有=,在加密本意就是占位符,在程序中没用, 网络传输中也浪费资源
        ### 所以需要重写base64加密去掉=, 传输回来验证时,需要在decode()方法之前把=补回来

        :param content: [传入需要加密的字符串]
        :type content: [str]
        :return: [返回base64加密, 把=去掉的加密]
        :rtype: [bin]
        """
        return base64.urlsafe_b64encode(content).replace(b'=', b'')

    @staticmethod
    def ba64decode(bit_str):
        """
        ### 把ba64encode()方法加密去掉的=补回来

        :param bit_str: [字节串]
        :type bit_str: [bin]
        :return: [返回=号补回来的base64加密字节串]
        :rtype: [bin]
        """
        # 补 =号的原理
        # base64长度能被4整除, 余0时不需要补=号, 其他情况都要根据余数的情况补=号.
        sem = len(bit_str) % 4  # 余数
        # 根据被除数4来判断余数是否补几个=好号
        if sem % 4 > 0:
            print('不能整除')
            bit_str += '=' * (4 - sem)
        return base64.urlsafe_b64decode(bit_str)

    # %% base64解密

    @staticmethod
    def decode(token, key):
        """
        ### 解密base64, 返回 payload 解密的结果

        :param token: [username, 加盐的钥匙]
        :type token: [bytes]
        :param key: [加盐的钥匙]
        :type key: [str]
        :return: [返回 payload 的结果]
        :rtype: [bytes]
        """
        # 拿.切割三个部分
        hander_bs, payload_bs, sign_bs = token.split(b'.')
        # 效验签名
        if isinstance(key, str):  # 如果是字符串的话,则进行转字节串
            key = key.encode()

        # 重新进行签名
        hm = hmac.new(key, hander_bs + b'.' +
                      payload_bs + b'.', digestmod='SHA256')

        # 对比传输过来的和重新签名的是否一致,否则不是本人, 如果一致就是本人,在进行下一步效验时间
        if sign_bs != Jwt.ba64encode(hm.digest()):
            print('对比结果不一致')
            raise  # 返回报错

        # 检查 exp是否过期
        payload_js = Jwt.ba64decode(payload_bs)

        # 这里需要注意,有的python版本传入字符串可能会报错
        payload = json.loads(payload_js)  # 转成字典,然后取值对比
        if 'exp' in payload:
            now = time.time()
            # 如果现在的时间大于签发的时间则时间到期, 到期提示报错
            if now > payload['exp']:
                print('签发时间过期')
                raise

        # 返回 payload部分的解码
        return payload


if __name__ == '__main__':
    # 传入用户名和加盐密匙, 用户名以字典形式
    token = Jwt.encode({'username': 'dd'}, '123456', 300)
    print(type(token), token)

    # 效验
    dc = Jwt.decode(token, '123456')
    print(dc)
